#ifndef _CUDA_MATRIX_CUH_
#define _CUDA_MATRIX_CUH_
/* common header for cuda matrix */
#include <cstdio>
#include "cuda_matrix.h"

namespace cuda_matrix{
    template<typename MT>
    inline void matfill_host( MT a, float f  ){
        for( int y = 0 ; y < a.num_line() ; y ++ ){
            float *e = (float*)((char*)a.elem + y * a.pitch);
            for( int x = 0 ; x < a.x_max ; x ++ ){
                e[ x ] = f;
            }
        }
    }    

    template<typename MT>
    inline void matcpy( MT dst, const MT src, enum cudaMemcpyKind kind ){
        cudaMemcpy2D( dst.elem , dst.pitch , src.elem, src.pitch , src.line_width() , src.num_line() , kind );
    }    

    template<typename MT>
    inline void matcpy_async( MT dst, const MT src, enum cudaMemcpyKind kind , cudaStream_t stream ){
        cudaMemcpy2DAsync( dst.elem , dst.pitch , src.elem, src.pitch , src.line_width() , src.num_line() , kind, stream  );
    }

    template<typename MT>
    inline cudaError_t mat_alloc_device( MT &m ){
        return cudaMallocPitch( (void**)&m.elem , &m.pitch , m.line_width(), m.num_line() );
    }

    template<typename MT>
    inline cudaError_t mat_free_device( MT m ){
        return cudaFree(m.elem);
    }

    template<typename MT>
    inline void mat_alloc_host( MT &m ){
        // align to 16
        m.pitch = ( ( ( m.x_max + 15 ) >> 4 ) << 4 ) * sizeof(float);
        m.elem  = (float*) malloc( m.num_bytes() );
    }

    template<typename MT>
    inline void mat_free_host( MT m ){
        free( m.elem );
    }

    template<>
    inline void matcpy<Mat4DBias>( Mat4DBias dst, const Mat4DBias src, enum cudaMemcpyKind kind ){
        matcpy<Mat4D>( dst.mat , src.mat , kind );
        cudaMemcpy( dst.v_bias , src.v_bias , src.mat.v_max * sizeof(float), kind ); 
        cudaMemcpy( dst.h_bias , src.h_bias , src.mat.h_max * sizeof(float), kind ); 
    }
    
    template<>
    inline cudaError_t mat_alloc_device<Mat4DBias>( Mat4DBias &m ){
        cudaError_t err = mat_alloc_device<Mat4D>( m.mat );
        cudaMalloc( (void **)&m.v_bias , m.mat.v_max * sizeof(float));
        cudaMalloc( (void **)&m.h_bias , m.mat.h_max * sizeof(float));
        return err;
    }

    template<>
    inline cudaError_t mat_free_device<Mat4DBias>( Mat4DBias m ){
        cudaError_t err = mat_free_device<Mat4D>( m.mat );        
        cudaFree( m.v_bias );
        cudaFree( m.h_bias );
        return err;
    }

    template<>
    inline void mat_alloc_host<Mat4DBias>( Mat4DBias &m ){
        mat_alloc_host<Mat4D>( m.mat );        
        m.v_bias = (float*) malloc( m.mat.v_max * sizeof(float) );
        m.h_bias = (float*) malloc( m.mat.h_max * sizeof(float) );
    }

    template<>
    inline void mat_free_host<Mat4DBias>( Mat4DBias m ){
        mat_free_host<Mat4D>( m.mat );        
        free( m.v_bias ); 
        free( m.h_bias );
    }
};

#include "cuda_conv2.cuh"

#endif

